from fontTools.misc.classifyTools import classify
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# 设置具体包含 config.json的目录

model_dir = r"D:\AI\model\Qwen3-4B-Thinking-2507" #这是一个文本生成模型

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(model_dir)
tokenizer = AutoTokenizer.from_pretrained(model_dir)


# 使用加载的模型和分词器创建生成文本的pipline
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)  # -1表示CPU

# 生成文本
# 方案1：使用max_new_tokens（推荐）
# text = text_generator(
#     "今天天气不错",
#     max_new_tokens=50,
#     num_return_sequences=1
# )

# 方案2：使用max_new_tokens（推荐）
text = text_generator(
    "今天天气不错",
    max_new_tokens=50, # 50个token,表示文本最多包含50个tokens
    num_return_sequences=1, # 参数num_return_sequences表示生成几个文本，值为1表示只生成一个文本
    truncation=True, # 启用截断 如果超出最大长度，则截断。如果是false，模型可能无法处理，可能会报错。
    temperature=0.7, # 控制文本的随机性，值越小越保守，值越大越随机越多样性，0.7是一个较为常用的参数
    top_k=50, # 采样的top_k个概率，这里表示只从概率最高的前50个中选择
    top_p=0.9, # 采样top_p概率，选择一组累计概率达到p的词汇，模型只会在这个概率区间内进行采样。0.9意味着模型会在可能性最强的90%的词中选下一个词。
    clean_up_tokenization_spaces=False #改参数控制生成的文本是否清理掉多余的空格，默认为True，生成文本会清除多余的空格
)

print(text)